#pragma once

#ifdef REFBASE_DEBUG
#include "log.h"
#define DEBUG_LOG(...) LOG(__VA_ARGS__)
#else // !REFBASE_DEBUG
#define DEBUG_LOG(...)
#endif

class RefBase;
class WeakRefCounter {
    friend class RefBase;

public:
    WeakRefCounter() = default;
    virtual ~WeakRefCounter() = default;

    void WeakRefCountIncrease();
    void WeakRefCountDecrease();
    RefBase *GetRefPtr() const;

private:
    int weakRefCount = 0;
    RefBase *obj = nullptr;
};

class RefBase {
public:
    RefBase();
    virtual ~RefBase() = default;

    void StrongRefCountIncrease();
    void StrongRefCountDecrease();
    WeakRefCounter *GetWeakRefCounter() const;

private:
    int strongRefCount = 0;
    WeakRefCounter *counter = nullptr;
};

template<class T>
class sptr {
public:
    sptr();
    sptr(T *obj);
    sptr(sptr<T> &&obj);
    sptr(const sptr<T> &obj);
    ~sptr();

    sptr<T> &operator=(T *obj);
    sptr<T> &operator=(sptr<T> &&obj);
    sptr<T> &operator=(const sptr<T> &obj);

    inline T *GetRefPtr()                      const { return pObject; }
    inline void ForceSetRefPtr(T *other)             { pObject = other; }

    inline T *operator->()                     const { return GetRefPtr(); }
    inline T &operator*()                      const { return *GetRefPtr(); }

    inline bool operator!=(const sptr<T> &obj) const { return obj.GetRefPtr() != GetRefPtr(); }
    inline bool operator==(const sptr<T> &obj) const { return obj.GetRefPtr() == GetRefPtr(); }
    inline bool operator!=(T *obj)             const { return obj != GetRefPtr(); }
    inline bool operator==(T *obj)             const { return obj == GetRefPtr(); }

private:
    T *pObject = nullptr;
};

template<class T>
class wptr {
public:
    wptr() = default;
    wptr(T *obj);
    wptr(const sptr<T> &ptr);
    ~wptr();

    wptr<T> &operator=(const sptr<T> &obj);
    wptr<T> &operator=(T *obj);

    sptr<T> promote();

    inline bool operator!=(const sptr<T> &obj) const { return obj.GetRefPtr() != counter->GetRefPtr(); }
    inline bool operator==(const sptr<T> &obj) const { return obj.GetRefPtr() == counter->GetRefPtr(); }
    inline bool operator!=(T *obj)             const { return obj != counter->GetRefPtr(); }
    inline bool operator==(T *obj)             const { return obj == counter->GetRefPtr(); }

private:
    WeakRefCounter *counter = nullptr;
};

// sptr implementation
template<class T>
sptr<T>::sptr()
{
    DEBUG_LOG("");
}

template<class T>
sptr<T>::sptr(T *obj)
{
    DEBUG_LOG("");
    pObject = obj;
    if (pObject != nullptr) {
        pObject->StrongRefCountIncrease();
    }
}

template<class T>
sptr<T>::sptr(sptr<T> &&obj)
{
    DEBUG_LOG("");
    pObject = obj.GetRefPtr();
    obj.ForceSetRefPtr(nullptr);
}

template<class T>
sptr<T>::sptr(const sptr<T> &obj)
{
    DEBUG_LOG("");
    pObject = obj.GetRefPtr();
    if (pObject != nullptr) {
        pObject->StrongRefCountIncrease();
    }
}

template<class T>
sptr<T>::~sptr()
{
    DEBUG_LOG("");
    if (pObject != nullptr) {
        pObject->StrongRefCountDecrease();
    }
}

template<class T>
sptr<T> &sptr<T>::operator=(sptr<T> &&obj)
{
    DEBUG_LOG("");
    if (pObject != nullptr) {
        pObject->StrongRefCountDecrease();
    }

    pObject = obj.GetRefPtr();
    obj.ForceSetRefPtr(nullptr);
    return *this;
}

template<class T>
sptr<T> &sptr<T>::operator=(const sptr<T> &obj)
{
    DEBUG_LOG("");
    if (pObject != nullptr) {
        pObject->StrongRefCountDecrease();
    }

    pObject = obj.GetRefPtr();
    if (pObject != nullptr) {
        pObject->StrongRefCountIncrease();
    }
    return *this;
}

template<class T>
sptr<T> &sptr<T>::operator=(T *obj)
{
    DEBUG_LOG("");
    if (pObject != nullptr) {
        pObject->StrongRefCountDecrease();
    }

    pObject = obj;
    if (pObject != nullptr) {
        pObject->StrongRefCountIncrease();
    }
    return *this;
}

// wptr implementation
template<class T>
wptr<T>::wptr(T *obj)
{
    if (obj != nullptr) {
        counter = obj->GetWeakRefCounter();
        counter->WeakRefCountIncrease();
    }
}

template<class T>
wptr<T>::wptr(const sptr<T> &ptr)
{
    if (ptr.GetRefPtr() != nullptr) {
        counter = ptr->GetWeakRefCounter();
        counter->WeakRefCountIncrease();
    }
}

template<class T>
wptr<T>::~wptr()
{
    if (counter != nullptr) {
        counter->WeakRefCountDecrease();
    }
}

template<class T>
wptr<T> &wptr<T>::operator=(const sptr<T> &obj)
{
    if (counter != nullptr) {
        counter->WeakRefCountDecrease();
    }

    if (obj.GetRefPtr() != nullptr) {
        counter = obj.GetRefPtr()->GetWeakRefCounter();
        counter->WeakRefCountIncrease();
    }
    return *this;
}

template<class T>
wptr<T> &wptr<T>::operator=(T *obj)
{
    if (counter != nullptr) {
        counter->WeakRefCountDecrease();
    }

    if (obj != nullptr) {
        counter = obj->GetWeakRefCounter();
        counter->WeakRefCountIncrease();
    }
    return *this;
}

template<class T>
sptr<T> wptr<T>::promote()
{
    if (counter != nullptr) {
        return reinterpret_cast<T *>(counter->GetRefPtr());
    }
    return nullptr;
}
